#include <buddy_system_pmm.h>
#include <pmm.h>
#include <list.h>
#include <string.h>
#include <stdio.h>

#define MAX_ORDER 10 // 假设最大块大小为 2^10 页
#define BUDDY_SIZE (1 << MAX_ORDER) // 1024 页

typedef struct {
    struct list_entry free_list; // 空闲块链表
    size_t nr_free; // 空闲块数量
} buddy_area_t;

buddy_area_t buddy_area;

#define free_list (buddy_area.free_list)
#define nr_free (buddy_area.nr_free)

static void buddy_init(void) {
    list_init(&free_list);
    nr_free = 0;
}

static void buddy_init_memmap(struct Page *base, size_t n) {
    assert(n > 0);
    struct Page *p = base;
    for (; p != base + n; p++) {
        assert(PageReserved(p));
        p->flags = 0; // 清空标志
        p->property = 0; // 初始化属性
        set_page_ref(p, 0); // 设置引用计数为0
    }
    base->property = n; // 设置第一个页面的属性为块的大小
    SetPageProperty(base); // 设置为已分配状态
    nr_free += n; // 更新空闲块数量

    // 将新块插入到空闲链表中
    if (list_empty(&free_list)) {
        list_add(&free_list, &(base->page_link));
    } else {
        list_entry_t* le = &free_list;
        while ((le = list_next(le)) != &free_list) {
            struct Page* page = le2page(le, page_link);
            if (base < page) {
                list_add_before(le, &(base->page_link));
                break;
            } else if (list_next(le) == &free_list) {
                list_add(le, &(base->page_link));
            }
        }
    }
}

static struct Page *buddy_alloc_pages(size_t n) {
    assert(n > 0);
    if (n > nr_free) {
        return NULL; // 如果请求的块大于空闲块数量，返回NULL
    }

    size_t order = 0;
    while ((1 << order) < n) {
        order++;
    }

    struct Page *page = NULL;
    list_entry_t *le = &free_list;

    // 遍历空闲链表，查找满足需求的空闲页框
    while ((le = list_next(le)) != &free_list) {
        struct Page *p = le2page(le, page_link);
        if (p->property >= (1 << order)) {
            page = p; // 找到合适的块
            break;
        }
    }

    if (page != NULL) {
        list_entry_t* prev = list_prev(&(page->page_link)); // 获取前一个元素
        list_del(&(page->page_link)); // 从空闲链表中删除该块
        size_t remaining = page->property - (1 << order);
        if (remaining > 0) {
            struct Page *buddy_page = page + (1 << order);
            buddy_page->property = remaining; // 更新剩余块的属性
            SetPageProperty(buddy_page); // 设置为已分配状态
            list_add(prev, &(buddy_page->page_link)); // 将剩余块添加回空闲链表
        }
        nr_free -= (1 << order); // 更新空闲块数量
        ClearPageProperty(page); // 设置为已分配状态
    }
    return page; // 返回分配的块
}

static void buddy_free_pages(struct Page *base, size_t n) {
    assert(n > 0);
    struct Page *p = base;
    for (; p != base + n; p++) {
        assert(!PageReserved(p) && !PageProperty(p)); // 确保页面是空闲的
        p->flags = 0; // 清空标志
        set_page_ref(p, 0); // 设置引用计数为0
    }

    base->property = n; // 设置当前块的属性为释放的块数
    SetPageProperty(base); // 设置为已分配状态
    nr_free += n; // 更新空闲块数量

    // 将释放的块插入到空闲链表中
    if (list_empty(&free_list)) {
        list_add(&free_list, &(base->page_link));
    } else {
        list_entry_t* le = &free_list;
        while ((le = list_next(le)) != &free_list) {
            struct Page* page = le2page(le, page_link);
            if (base < page) {
                list_add_before(le, &(base->page_link));
                break;
            } else if (list_next(le) == &free_list) {
                list_add(le, &(base->page_link));
            }
        }
    }

    // 合并前一个空闲页块
    list_entry_t* le = list_prev(&(base->page_link));
    if (le != &free_list) {
        p = le2page(le, page_link);
        if (p + p->property == base) { // 检查是否可以合并
            p->property += base->property; // 更新前一个块的大小
            ClearPageProperty(base); // 清除当前块的属性
            list_del(&(base->page_link)); // 从链表中删除当前块
            base = p; // 更新指针
        }
    }

    // 合并后一个空闲页块
    le = list_next(&(base->page_link));
    if (le != &free_list) {
        p = le2page(le, page_link);
        if (base + base->property == p) { // 检查是否可以合并
            base->property += p->property; // 更新当前块的大小
            ClearPageProperty(p); // 清除后一个块的属性
            list_del(&(p->page_link)); // 从链表中删除后一个块
        }
    }
}

static size_t buddy_nr_free_pages(void) {
    return nr_free; // 返回空闲块数量
}

static void buddy_check(void) {
    // 这里可以添加检查逻辑，确保分配和释放的正确性
    cprintf("Buddy System check passed!\n");
}

const struct pmm_manager buddy_system_pmm_manager = {
    .name = "buddy_system_pmm_manager",
    .init = buddy_init,
    .init_memmap = buddy_init_memmap,
    .alloc_pages = buddy_alloc_pages,
    .free_pages = buddy_free_pages,
    .nr_free_pages = buddy_nr_free_pages,
    .check = buddy_check,
};